Repository navigation
Add GPU support for Cellpose - #39
Conversation
|
Warning Newer version of the nf-core template is available. Your pipeline is using an old version of the nf-core template: 4.0.3. For more documentation on how to update your pipeline, please see the Synchronisation documentation. |
|
quentinblampey
left a comment
quentinblampey
left a comment
There was a problem hiding this comment.
Thanks @alihamraoui, I made some comments (it's mostly questions actually)!
| container "${ workflow.containerEngine in ['singularity', 'apptainer'] && !task.ext.singularity_pull_docker_container | ||
| ? 'https://community-cr-prod.seqera.io/docker/registry/v2/blobs/sha256/22/22d62d6425b70620138ad8764139528c1acaabf6cd06403134c8439caa1c9a31/data' | ||
| : 'community.wave.seqera.io/library/python_sopa_cellpose:d098579826bbcf24' }" | ||
| // CPU image: cellpose v3 | GPU image (task.ext.use_gpu): cellpose v4 + pytorch/CUDA, built from patch_segmentation_cellpose/environment_gpu.yml |
There was a problem hiding this comment.
We don't need a GPU to download a model, we can use the original CPU image
There was a problem hiding this comment.
Only PATCH_SEGMENTATION_CELLPOSE has the process_gpu label, this process doesn't use the GPU for the download, it runs on the CPU. it just reuses the v4 image because sopa download cellpose caches the model for the installed Cellpose version, and the CPU image (v3) has no cpsam model.
There was a problem hiding this comment.
Oh I see, makes sense
| container "${ workflow.containerEngine in ['singularity', 'apptainer'] && !task.ext.singularity_pull_docker_container | ||
| ? 'https://community-cr-prod.seqera.io/docker/registry/v2/blobs/sha256/22/22d62d6425b70620138ad8764139528c1acaabf6cd06403134c8439caa1c9a31/data' | ||
| : 'community.wave.seqera.io/library/python_sopa_cellpose:d098579826bbcf24' }" | ||
| // CPU image: cellpose v3 | GPU image (task.ext.use_gpu): cellpose v4 + pytorch/CUDA, built from patch_segmentation_cellpose/environment_gpu.yml |
There was a problem hiding this comment.
Same here: we don't need a GPU to resolve shapes, we can revert the changes
There was a problem hiding this comment.
I think since the GPU image is already pulled for the patches, reusing the same image avoids an extra ~4 GB pull. what do you think?
There was a problem hiding this comment.
Yes, true, good point!
| ext.use_gpu = params.cellpose_use_gpu | ||
| } | ||
| withName: PATCH_SEGMENTATION_CELLPOSE { | ||
| accelerator = { params.cellpose_use_gpu ? 1 : null } |
There was a problem hiding this comment.
What's the accelerator doing?
There was a problem hiding this comment.
accelerator = 1 allocate one GPU to the task, I think by adding --gres=gpu:1 to the run command for singularity for example. it's only applied to PATCH_SEGMENTATION_CELLPOSE right now.
There was a problem hiding this comment.
Is nextflow not doing it by itself? We really need to set this ourselves?
There was a problem hiding this comment.
I'm not sure accelerator = 1 alone is enough. According to the docs, it declares the number of GPUs a task needs, but only some executors use it (on Slurm, clusterOptions '--gres=gpu:1' may still be needed). The container also needs the driver, with --nv for Singularity or --gpus all for Docker, which is what -profile gpu adds. I'm running some tests to make sure everything works.
There was a problem hiding this comment.
Yeah it may be right, I'm just surprised that nextflow doesn't already handle all this stuff itself by default, but maybe you're right
| : 'community.wave.seqera.io/library/python_sopa_cellpose:d098579826bbcf24' }" | ||
| // CPU image: cellpose v3 | GPU image (task.ext.use_gpu): cellpose v4 + pytorch/CUDA, built from patch_segmentation_cellpose/environment_gpu.yml | ||
| container "${ task.ext.use_gpu | ||
| ? 'community.wave.seqera.io/library/python_sopa_cellpose_pytorch-gpu_cuda-version:7ca3fb7cbc5a048b' |
There was a problem hiding this comment.
Nice! Did you test this container on a machine with CUDA? Does it automatically detect the existing CUDA drivers?
❌ nf-test failed with latest Nextflow versionNote Tests with Nextflow's latest version failed but it will not cause a CI workflow failure.
See the full run for details. |
|
I tested the whole pipe on a machine with CUDA (cuda 12) with different param combinations.
|
quentinblampey
left a comment
quentinblampey
left a comment
There was a problem hiding this comment.
I added some comments again, we're almost there :)
| } | ||
| def cellpose_v4 = params.cellpose_version != null ? params.cellpose_version.toString() == '4' : workflow.profile.contains('gpu') | ||
| if (params.use_cellpose && cellpose_v4 && params.cellpose_model_type != null) { | ||
| log.warn("Cellpose v4 only provides the 'cpsam' model: 'cellpose_model_type=${params.cellpose_model_type}' will be ignored.") |
There was a problem hiding this comment.
You mean Cellpose v4 only supports the pretrained_model argument, no? As far as I remember, we can provided multiple pretrained_model, and cpsam is just one of them, but indeed model_type is only for v3
There was a problem hiding this comment.
You are right, v4 now supports multiple pretrained models, not just cpsam.
I checked, the other pretrained models are available since v4.2.1, but we're using 4.1. I assume they are also compatible with 4.1, so I'll remove the warning.
There was a problem hiding this comment.
I think it's better to build a new image with Cellpose >=4.2.1
| log.warn("Cellpose v4 only provides the 'cpsam' model: 'cellpose_model_type=${params.cellpose_model_type}' will be ignored.") | ||
| } | ||
| if (params.use_cellpose && !cellpose_v4 && workflow.profile.contains('gpu')) { | ||
| log.warn("'cellpose_version=3' is used with the 'gpu' profile, but the Cellpose v3 image is not built with CUDA: Cellpose may fall back to CPU. Use 'cellpose_version=4' for GPU support.") |
There was a problem hiding this comment.
Good point. Do we need another image for cellpose v3 + GPU then?
There was a problem hiding this comment.
or, make a single v3 image that works for both cases: v3 + GPU and v3 on CPU. It will be a larger image for CPU-only users, but it keeps the module simpler. What do you think?
There was a problem hiding this comment.
Yes, it's also possible, as you prefer (if it's not a much much bigger image)
| ? 'https://community-cr-prod.seqera.io/docker/registry/v2/blobs/sha256/22/22d62d6425b70620138ad8764139528c1acaabf6cd06403134c8439caa1c9a31/data' | ||
| : 'community.wave.seqera.io/library/python_sopa_cellpose:d098579826bbcf24' }" | ||
|
|
||
| container "${ task.ext.cellpose_v4 |
There was a problem hiding this comment.
I'm wondering, is it possible to move this big container logic within the modules.config? E.g., within the withName: 'DOWNLOAD_CELLPOSE_MODEL|PATCH_SEGMENTATION_CELLPOSE|RESOLVE_CELLPOSE'?
I don't know, maybe it's not possible, just an idea
There was a problem hiding this comment.
Good point, yes, it's possible. It even avoids repeated code.
I see also that -profile conda never gets v4 (it always uses environment.yml). I'll handle it.
|
Hi @quentinblampey, As expected, Cellpose v4.1.1 only accepts Now we have 4 images: v3 and v4, each with its GPU (pytorch/CUDA) version. Sopa now uses the Cellpose v3 image by default, and Tested with profile |
quentinblampey
left a comment
quentinblampey
left a comment
There was a problem hiding this comment.
Hi @alihamraoui, thanks for the massive work, this is a really great improvement!! I just made a minor comment, I'm not sure if this has to be removed or not. I approved anyway, I'll let you see if it needs to be corrected or not and merge it :)
|
|
||
| conda "${moduleDir}/environment.yml" | ||
|
|
||
| container "${ workflow.containerEngine in ['singularity', 'apptainer'] && !task.ext.singularity_pull_docker_container |
There was a problem hiding this comment.
Do we still need the container command within these process, since we have it already in conf/modules.config?
There was a problem hiding this comment.
I keept them as a fallback for module tests or nf-core download. conf/modules.config overrides them to pick the GPU or v4 variant. but we can remove them and nothing changes when the pipeline runs.
|
Good! thanks @quentinblampey, |
2 participants
--cellpose_use_gpuswitches to a Cellpose v4 + PyTorch/CUDA image.New warnings when:
cellpose_use_gpuis set without-profile gpu(the container would not see the GPU);cellpose_model_typeis set in GPU mode (Cellpose v4 only shipscpsam, so the value is ignored).Usage
Tested on a GPU node with Singularity